package Tree.Trie;

import java.util.HashMap;
import java.util.Map;

/**
 * Lettcode 677: 键值映射
 * 实现一个 MapSum 类里的两个方法，insert 和 sum。
 *
 * 对于方法 insert，你将得到一对（字符串，整数）的键值对。字符串表示键，整数表示值。如果键已经存在，那么原来的键值对将被替代成新的键值对。
 *
 * 对于方法 sum，你将得到一个表示前缀的字符串，你需要返回所有以该前缀开头的键的值的总和。
 *
 * 示例 1:
 *
 * 输入: insert("apple", 3), 输出: Null
 * 输入: sum("ap"), 输出: 3
 * 输入: insert("app", 2), 输出: Null
 * 输入: sum("ap"), 输出: 5
 *
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode-cn.com/problems/map-sum-pairs
 * 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
 */

/**
 * 方法一：利用map进行存储
 */
public class MapSum {
    private Map<String, Integer> stringIntegerMap;

    /** Initialize your data structure here. */
    public MapSum() {
        stringIntegerMap = new HashMap<>();
    }

    public void insert(String key, int val) {
        stringIntegerMap.put(key, val);
    }

    public int sum(String prefix) {
        int sum = 0;
        for (String string : stringIntegerMap.keySet()) {
            if (string.startsWith(prefix)) {
                sum += stringIntegerMap.get(string);
            }
        }
        return sum;
    }
}

/**
 * 方法二： 前缀树
 */
class MapSumI{
    private class Node {
        Node[] childs = new Node[26];
        int value;
    }

    Node root = new Node();

    /** Initialize your data structure here. */
    public MapSumI() {

    }

    public void insert(String key, int val) {
        insert(key, val, root);
    }

    private void insert(String key, int val, Node node) {
        if (node == null) {
            return;
        }
        if (key.length() == 0) {
            node.value = val;
            return;
        }
        int index = indexForChar(key.charAt(0));
        if (node.childs[index] == null) {
            node.childs[index] = new Node();
        }
        insert(key.substring(1), val, node.childs[index]);
    }

    public int sum(String prefix) {
        return sum(prefix, root);
    }

    private int sum(String prefix, Node node) {
        if (node == null) {
            return 0;
        }
        if (prefix.length() != 0) {
            int index = indexForChar(prefix.charAt(0));
            return sum(prefix.substring(1), node.childs[index]);
        }
        int sum = node.value;
        for (Node child: node.childs) {
            sum += sum(prefix, child);
        }
        return sum;
    }

    private int indexForChar(char c) {
        return c - 'a';
    }
}
